import subprocess
import os
import argparse
import re
import sys
import time
import json
import tempfile
from tqdm import tqdm
from pydub import AudioSegment
from pydub.silence import detect_silence

def detect_silent_segments(input_mkv, silence_thresh=-50, min_silence_len=500):
    """检测视频中的静音片段"""
    try:
        # 创建临时文件路径
        temp_audio_path = tempfile.mktemp(suffix='.wav')
        
        print(f"提取音频到临时文件: {temp_audio_path}")
        cmd = [
            'ffmpeg', '-y', 
            '-i', input_mkv,
            '-vn', '-acodec', 'pcm_s16le', '-ar', '44100', '-ac', '1',
            temp_audio_path
        ]
        
        process = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
        
        # 等待处理完成
        while process.poll() is None:
            time.sleep(0.1)
        
        if process.returncode != 0:
            print(f"音频提取失败 (返回码 {process.returncode}):")
            print(process.stderr.read())
            return []
        
        print("加载音频并检测静音...")
        audio = AudioSegment.from_wav(temp_audio_path)
        silence_ranges = detect_silence(
            audio, min_silence_len=min_silence_len, silence_thresh=silence_thresh
        )
        
        # 清理临时文件
        if os.path.exists(temp_audio_path):
            os.unlink(temp_audio_path)
            
        print(f"检测到 {len(silence_ranges)} 个静音片段")
        return silence_ranges
    except Exception as e:
        print(f"提取音频失败: {e}")
        return []

def get_video_duration(input_file):
    """获取视频时长（秒）"""
    try:
        print("获取视频时长...")
        # 使用ffprobe获取视频时长
        cmd = ['ffprobe', '-v', 'error', '-show_entries', 
               'format=duration', '-of', 'json', input_file]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        # 检查是否成功
        if result.returncode == 0:
            data = json.loads(result.stdout)
            duration = float(data['format']['duration'])
            print(f"视频时长: {duration:.2f}秒")
            return duration
        
        print(f"ffprobe错误 (返回码 {result.returncode}):")
        print(result.stderr)
        return 0
    except Exception as e:
        print(f"获取视频时长失败: {e}")
        return 0

def get_total_frames(input_file):
    """获取视频总帧数（使用更可靠的方法）"""
    try:
        print("尝试获取视频总帧数...")
        # 方法1：使用ffprobe获取视频时长和帧率
        cmd = ['ffprobe', '-v', 'error', 
               '-select_streams', 'v:0',
               '-show_entries', 'stream=avg_frame_rate,duration', 
               '-of', 'json', input_file]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode == 0:
            data = json.loads(result.stdout)
            streams = data.get('streams', [])
            if streams:
                stream = streams[0]
                duration = float(stream.get('duration', 0))
                frame_rate_str = stream.get('avg_frame_rate', '0/0')
                
                # 解析帧率 (如 "30000/1001")
                if '/' in frame_rate_str:
                    num, den = frame_rate_str.split('/')
                    frame_rate = float(num) / float(den) if den != '0' else 0
                else:
                    frame_rate = float(frame_rate_str) if frame_rate_str else 0
                
                if duration > 0 and frame_rate > 0:
                    total_frames = int(duration * frame_rate)
                    print(f"通过时长和帧率计算总帧数: {total_frames}")
                    return total_frames
        
        # 方法2：使用ffmpeg获取帧数
        print("尝试通过ffmpeg获取帧数...")
        cmd = ['ffmpeg', '-i', input_file, '-map', '0:v:0', '-c', 'copy', '-f', 'null', '-']
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        # 从输出中解析帧数
        frame_match = re.search(r'frame=\s*(\d+)', result.stderr)
        if frame_match:
            total_frames = int(frame_match.group(1))
            print(f"通过ffmpeg输出获取总帧数: {total_frames}")
            return total_frames
        
        print("无法获取精确总帧数")
        return 0
    except Exception as e:
        print(f"获取总帧数失败: {e}")
        return 0

def generate_filter_script(silence_ranges, duration):
    """生成FFmpeg复杂滤镜脚本"""
    if not silence_ranges:
        print("没有静音片段需要移除")
        return None

    # 将静音时间转换为秒
    silence_segments = [(start/1000.0, end/1000.0) for start, end in silence_ranges]
    
    # 构建裁剪点列表
    cut_points = [0.0]
    for start, end in silence_segments:
        cut_points.append(start)
        cut_points.append(end)
    cut_points.append(duration)
    
    # 生成视频分段
    segments = []
    for i in range(0, len(cut_points), 2):
        if i+1 < len(cut_points):
            start = cut_points[i]
            end = cut_points[i+1]
            segments.append((start, end))
    
    # 构建FFmpeg滤镜链
    filter_script = ""
    concat_inputs = []
    for idx, (start, end) in enumerate(segments):
        filter_script += (
            f"[0:v]trim=start={start}:end={end},setpts=PTS-STARTPTS[v{idx}];\n"
            f"[0:a]atrim=start={start}:end={end},asetpts=PTS-STARTPTS[a{idx}];\n"
        )
        concat_inputs.append(f"[v{idx}]")
        concat_inputs.append(f"[a{idx}]")
    
    filter_script += f"{''.join(concat_inputs)}concat=n={len(segments)}:v=1:a=1[outv][outa]"
    
    print(f"生成滤镜脚本，包含 {len(segments)} 个视频片段")
    return filter_script

def run_ffmpeg_with_progress(cmd, total_frames):
    """运行FFmpeg并显示进度条"""
    # 创建进度条
    progress_bar = tqdm(total=total_frames, unit='帧', desc='处理进度', disable=(total_frames <= 0))
    
    # 运行FFmpeg
    process = subprocess.Popen(cmd, stderr=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
    
    # 进度跟踪
    frame_pattern = re.compile(r'frame=\s*(\d+)')
    time_pattern = re.compile(r'time=\s*(\d+:\d+:\d+\.\d+)')
    last_frame = 0
    
    while True:
        line = process.stderr.readline()
        if not line:
            break
            
        # 更新进度
        frame_match = frame_pattern.search(line)
        time_match = time_pattern.search(line)
        
        if frame_match:
            current_frame = int(frame_match.group(1))
            if total_frames > 0:
                progress_bar.update(current_frame - last_frame)
            else:
                progress_bar.set_description(f"处理进度 (帧: {current_frame})")
            last_frame = current_frame
        elif time_match:
            current_time = time_match.group(1)
            progress_bar.set_description(f"处理进度 (时间: {current_time})")
    
    # 等待进程结束
    return_code = process.wait()
    progress_bar.close()
    
    return return_code, process.stdout.read(), process.stderr.read()

def main():
    parser = argparse.ArgumentParser(description='移除视频静音部分并输出为MP4')
    parser.add_argument('input', help='输入MKV视频文件路径')
    parser.add_argument('output', help='输出MP4视频文件路径')
    parser.add_argument('--silence_thresh', type=int, default=-50, 
                        help='静音阈值(dBFS)，默认-50')
    parser.add_argument('--min_silence_len', type=int, default=500,
                        help='最小静音时长(ms)，默认500')
    parser.add_argument('--gpu', action='store_true',
                        help='启用GPU加速（需要支持CUDA的GPU）')
    args = parser.parse_args()

    # 验证输入文件
    if not os.path.exists(args.input):
        print(f"错误: 输入文件不存在: {args.input}")
        sys.exit(1)
    
    print(f"处理文件: {args.input}")
    
    # 获取视频时长
    duration = get_video_duration(args.input)
    if duration <= 0:
        print(f"警告: 无法获取视频时长，使用默认值60秒")
        duration = 60  # 默认值
    
    # 获取总帧数（用于进度条）
    total_frames = get_total_frames(args.input)
    if total_frames <= 0:
        print("警告: 无法获取精确总帧数，进度条可能不准确")
    
    # 检测静音片段
    silence_ranges = detect_silent_segments(
        args.input, args.silence_thresh, args.min_silence_len
    )
    
    if not silence_ranges:
        print("未检测到静音片段，将直接转换格式")
    
    # 生成滤镜脚本
    filter_script = generate_filter_script(silence_ranges, duration)
    
    # 构建FFmpeg命令
    if filter_script:
        print("正在处理视频并移除静音部分...")
        cmd = [
            'ffmpeg', '-y', 
            '-i', args.input,
            '-filter_complex', filter_script,
            '-map', '[outv]', '-map', '[outa]',
            '-c:v', 'libx264', '-preset', 'fast',
            '-c:a', 'aac', '-b:a', '192k',
            args.output
        ]
    else:
        print("正在转换视频格式...")
        # 没有静音部分时直接转换格式
        cmd = [
            'ffmpeg', '-y', 
            '-i', args.input,
            '-c:v', 'copy', 
            '-c:a', 'aac', '-b:a', '192k',
            args.output
        ]
    
    # 如果启用GPU加速
    if args.gpu:
        print("启用GPU加速...")
        # 在输入前添加GPU加速参数
        cmd = cmd[:1] + ['-hwaccel', 'cuda', '-hwaccel_output_format', 'cuda'] + cmd[1:]
        
        # 替换视频编码器为GPU版本
        if filter_script:
            cmd[cmd.index('-c:v')+1] = 'h264_nvenc'
            # 添加GPU编码参数
            cmd.insert(cmd.index('-c:v')+2, '-preset')
            cmd.insert(cmd.index('-c:v')+3, 'p7')
            cmd.insert(cmd.index('-c:v')+4, '-tune')
            cmd.insert(cmd.index('-c:v')+5, 'hq')
            cmd.insert(cmd.index('-c:v')+6, '-cq')
            cmd.insert(cmd.index('-c:v')+7, '18')
            cmd.insert(cmd.index('-c:v')+8, '-rc')
            cmd.insert(cmd.index('-c:v')+9, 'vbr')
    
    # 打印完整的FFmpeg命令
    print("执行命令:")
    print(" ".join(cmd))
    
    try:
        # 运行FFmpeg并显示进度
        return_code, stdout, stderr = run_ffmpeg_with_progress(cmd, total_frames)
        
        if return_code != 0:
            print(f"FFmpeg处理失败 (返回码 {return_code}):")
            print(stderr)
            sys.exit(1)
        
        print(f"\n处理完成！输出文件: {args.output}")
    except Exception as e:
        print(f"处理失败: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()